import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from datetime import datetime
from models.detection_model import DetectionModel, ModelVersion
from database.db import Base

# 数据库连接
DATABASE_URL = "sqlite:///./pilot_manage.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建数据库表
Base.metadata.create_all(bind=engine)

# 检测模型示例数据
detection_models_data = [
    {
        "name": "管道泄漏检测模型",
        "description": "用于检测管道泄漏的YOLOv8模型",
        "format": "YOLOv8",
        "application_scenario": "管道泄漏检测，适用于石油、天然气、水务等管道系统的泄漏监测",
        "input_requirements": "图片分辨率：1920x1080，视频：1080p 30fps",
        "output_format": "边界框坐标(x1,y1,x2,y2) + 置信度 + 泄漏类型",
        "versions": [
            {
                "version_name": "V1.0基础版",
                "version_code": "1.0.0",
                "accuracy": "85%",
                "recall": "80%",
                "description": "初始版本，支持基本的管道泄漏检测"
            },
            {
                "version_name": "V1.1优化版",
                "version_code": "1.1.0",
                "accuracy": "92%",
                "recall": "88%",
                "description": "优化了小泄漏点检测能力，提高了准确率"
            }
        ]
    },
    {
        "name": "电力线破损检测模型",
        "description": "基于Mask R-CNN的电力线破损检测模型",
        "format": "Mask R-CNN",
        "application_scenario": "电力线、电线杆破损检测，适用于电力巡检",
        "input_requirements": "图片分辨率：2560x1440，支持倾斜摄影图像",
        "output_format": "实例分割掩码 + 破损类型 + 置信度",
        "versions": [
            {
                "version_name": "V1.0初版",
                "version_code": "1.0.0",
                "accuracy": "88%",
                "recall": "82%",
                "description": "初始版本，支持电力线断裂、绝缘子破损检测"
            }
        ]
    },
    {
        "name": "植被覆盖分析模型",
        "description": "用于分析植被覆盖情况的语义分割模型",
        "format": "SegFormer",
        "application_scenario": "农业监测、生态环境评估、城市绿化分析",
        "input_requirements": "多光谱遥感图像，分辨率不低于0.5m/pixel",
        "output_format": "语义分割图 + 植被覆盖率统计",
        "versions": [
            {
                "version_name": "V1.0标准版",
                "version_code": "1.0.0",
                "accuracy": "95%",
                "recall": "93%",
                "description": "标准版，支持主要植被类型识别"
            },
            {
                "version_name": "V2.0增强版",
                "version_code": "2.0.0",
                "accuracy": "97%",
                "recall": "96%",
                "description": "增强版，增加了季节性植被变化监测能力"
            },
            {
                "version_name": "V2.1农业专用版",
                "version_code": "2.1.0",
                "accuracy": "98%",
                "recall": "97%",
                "description": "农业专用版，优化了农作物类型识别"
            }
        ]
    }
]

def init_detection_models():
    """初始化检测模型数据"""
    db = SessionLocal()
    try:
        print("开始初始化检测模型数据...")
        
        # 检查是否已有数据
        existing_count = db.query(DetectionModel).count()
        if existing_count > 0:
            print(f"检测到已有{existing_count}个检测模型，跳过初始化")
            return
        
        # 创建检测模型和版本
        for model_data in detection_models_data:
            # 创建检测模型
            model = DetectionModel(
                name=model_data["name"],
                description=model_data["description"],
                format=model_data["format"],
                application_scenario=model_data["application_scenario"],
                input_requirements=model_data["input_requirements"],
                output_format=model_data["output_format"],
                created_at=datetime.now(),
                updated_at=datetime.now()
            )
            db.add(model)
            db.flush()  # 获取model.id
            
            # 创建版本
            versions = []
            for version_data in model_data["versions"]:
                version = ModelVersion(
                    model_id=model.id,
                    version_name=version_data["version_name"],
                    version_code=version_data["version_code"],
                    accuracy=version_data["accuracy"],
                    recall=version_data["recall"],
                    description=version_data["description"],
                    created_at=datetime.now(),
                    updated_at=datetime.now()
                )
                db.add(version)
                db.flush()
                versions.append(version)
            
            # 设置最新版本为当前版本
            if versions:
                model.current_version_id = versions[-1].id  # 最后一个版本为最新版本
        
        db.commit()
        print(f"成功初始化{len(detection_models_data)}个检测模型")
        
    except Exception as e:
        print(f"初始化检测模型数据失败: {str(e)}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == "__main__":
    init_detection_models()